#include <stdio.h>
#include <gmp.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

mpz_t p, q;  //RSA算法中中途生成的两个素数
mpz_t n;     //n = p*q
mpz_t tmp_N; // tmp_N = (p-1)*(q-1)
mpz_t e;     //公钥为(e,n)
mpz_t d;     //私钥为(d,n)

gmp_randstate_t greatRnumer; //随机生成的大数

mpz_t M, C;  //M为明文，C是加密结果
mpz_t tmp_M; //tmp_M用来存放解密之后的明文，用来和加密的明文进行对比

char *Message;        // 明文
char *tmp_Message;    //解密后的明文字符串
char* tmp_EM;   //用来存放解密后的EM
char *PS;             //中间随机生成的一定长度的字符串
char *EM;             //实际加密使用的字符串
char *Ciphertext;     //密文
int mLen = 0;         //明文的长度
unsigned long long k; //长度限制，n的字节数

/**
 * 生成 p, q, n, phi(n), e 等参数
 * 通过循环进行生成合适的密钥
 * @param bit int 要求的 k 的位数
*/
void generateKey(int bit);

/**
 * 生成PS字符串
 * @param mlen int 明文的长度
*/
void getPS(int mlen);

/**
 * 得到处理之后的实际在加密中使用的字符串EM
*/
void getEM();

/**
 * 初始化所有的字符串的空间，以及k的长度
 */
void initString();

/**
 * 最后释放所有的大数以及字符串所占有的空间
*/
void clearAll();
/**
 * 将字符串转换为大数
 * @param res mpz_t 目标大数
 * @param src char* 源字符串
 * @param length int 长度，一般用 k 来作为参数
*/
void OS2IP(mpz_t res, char src[], unsigned long length);

/**
 * 将大数装换为字符串
 * @param src mpz_t 源大数
 * @param dst char* 目标字符串
*/
void I2OSP(char res[], mpz_t src);

/**
 * 解密函数
*/
void decode();

/**
 * 加密函数
*/
void encode();

int main()
{
    //首先生成密钥
    generateKey(1024);

    initString();

    printf("The plaintext is %s\n\n", Message);

    encode();

    decode();

    printf("After decoding, the final plaintext is %s\n\n", tmp_Message);

    if (!strcmp(Message, tmp_Message))
    {
        printf("The RSA is right!\n");
    }

    clearAll();
}

void getPS(int mLen)
{
    int i;
    PS = (char *)malloc(k - mLen - 2);
    PS[k - mLen - 3] = 0;
    for (i = 0; i < k - mLen - 3; i++)
    {
        PS[i] = rand() % 255 + 1;
    }
}

void getEM()
{
    int i;
    //设置EM字符串各个部分的长度
    int p1 = 1, p2 = 1, p3 = k - mLen - 3, p4 = 1, p5 = mLen;
    EM[0] = 0;
    EM[1] = 2;
    for (i = 0; i < p3; i++)
    {
        EM[p1 + p2 + i] = PS[i];
    }
    EM[p1 + p2 + p3] = 0;
    for (i = 0; i < p5; i++)
    {
        EM[p1 + p2 + p3 + p4 + i] = Message[i];
    }
}

void generateKey(int bit)
{
    while (1)
    {
        //随机生成大数
        gmp_randinit_default(greatRnumer);
        gmp_randseed_ui(greatRnumer, time(NULL));

        //初始化素数p,q
        mpz_init(p);
        mpz_init(q);

        //随机生成两个大数
        mpz_urandomb(p, greatRnumer, (bit + 1) / 2);
        mpz_urandomb(q, greatRnumer, (bit + 1) / 2);

        //生成素数
        mpz_nextprime(p, p);
        mpz_nextprime(q, q);

        //生成大数n = p*q
        mpz_init(n);
        mpz_mul(n, p, q);

        if (mpz_sizeinbase(n, 2) == bit)
        { //用来判断是否生成的是合适的位数
            break;
        }
    }
    //计算n的欧拉函数的值
    mpz_init(tmp_N);
    mpz_sub_ui(p, p, 1);
    mpz_sub_ui(q, q, 1);
    mpz_mul(tmp_N, p, q);

    //输出p,q的长度
    gmp_printf("p: %d q: %d n: %d\n\n", mpz_sizeinbase(p, 2), mpz_sizeinbase(q, 2), mpz_sizeinbase(n, 2));

    //公钥
    //e通常使用的是3,17,65537
    mpz_init_set_ui(e, 65537);
    gmp_printf("Public key is (%ZX, %ZX)\n\n", e, n);

    //私钥 d为e mod tmp_N 的逆元
    mpz_init(d);
    mpz_invert(d, e, tmp_N);
    gmp_printf("Private key is (%ZX, %ZX)\n\n", d, n);
}

void initString()
{
    //首先是确定明文
    int i;
    Message = "I love SYSU !";
    mLen = strlen(Message);
    unsigned long long size = mpz_sizeinbase(n, 2);
    k = size / 8 + (size % 8 ? 1 : 0);

    //其他字符串都要分配足够长的空间
    Ciphertext = (char *)malloc(k + 1);
    EM = (char *)malloc(k + 1);
    tmp_EM = (char *)malloc(k+1);
    tmp_Message = (char *)malloc(k + 1);
    for (int i = 0; i < k; i++)
    {
        tmp_EM[i] = Ciphertext[i] = tmp_Message[i] = EM[i] = 0;
    }
}

void clearAll()
{
    mpz_clear(d);
    mpz_clear(e);
    mpz_clear(n);
    mpz_clear(p);
    mpz_clear(q);
    mpz_clear(tmp_N);
    mpz_clear(C);
    mpz_clear(M);
    mpz_clear(tmp_M);

    free(tmp_Message);
    free(tmp_EM);
    free(PS);
    free(EM);
    free(Ciphertext);
}

void OS2IP(mpz_t res, char src[], unsigned long len)
{
    int i;
    mpz_init(res);
    for (i = len - 1; i >= 0; i--)
    {
        mpz_mul_ui(res, res, 256);
        mpz_add_ui(res, res, src[i] & 0x0000ff);
    }
}

void I2OSP(char res[], mpz_t src)
{
  unsigned long long size = mpz_sizeinbase(src, 2);
  int xLen = size / 8 + (size % 8 ? 1 : 0);
  free(res);
  res = (char *)malloc(xLen + 1);
  mpz_t tmp, copy;
  mpz_init_set(copy, src);

  // 循环生成字符串
  for (int i = 0; i < xLen; i++)
  {
    mpz_init(tmp);
    mpz_mod_ui(tmp, copy, 256);
    res[i] = mpz_get_ui(tmp) & 0x0000ff;
    mpz_div_ui(copy, copy, 256);
  }
}

void encode(){
    //首先先判断长度的问题
    if(k - mLen <11){
        printf("The message is too long");
        return ;
    }
    //处理字符串
    getPS(mLen);
    getEM();
    OS2IP(M,EM,k);
    gmp_printf("Encode > the message(mpz_t) is %ZX\n\n",M);
    mpz_powm(C,M,e,n);
    I2OSP(Ciphertext,C);
    // printf("Encode > The Cipertext is %s\n\n",Ciphertext);
}

void decode(){
    //判断传入的n的字节数和k是否一致，如果不一致，则输出错误
    unsigned long long size = mpz_sizeinbase(n, 2);
    unsigned long long n_byte = size / 8 + (size % 8 ? 1 : 0);
    if(n_byte!=k || k<11){
        printf("decryption error\n\n");
        return;
    }

    //开始解密
    OS2IP(C,Ciphertext,k);
    mpz_init(tmp_M);
    mpz_powm(tmp_M,C,d,n);
    gmp_printf("Decode > the final message(mpz_t) is %ZX\n\n",tmp_M);
    
    
    I2OSP(tmp_EM,tmp_M);

    int start;//用来判断Message在EM中的位置
    int i; //循环变量
    for(start = 1;start<k;start++){
        if(tmp_EM[start]==0){
            break;
        }
    }
    start = start+1;

    for(i=0;i<(k - start);i++){
        tmp_Message[i] = tmp_EM[start+i];
    }
}





